# File: EncoderCompensationPlot.py
# Created by: JH
# Created on: 20191010
# Last Update by: JH
# Last Update: 20200513
# Version: 2.02
# Coding: utf-8


import os
import re
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import sys

begin_row = 0
end_row = 990000
isrtick_col = 0
heidenhain_col = 4
icmu_col = 5
tw8_col = 6
present_mode_col = 1
index_current_col = 3

standard_deg = 0.0515 #150 count for icmu, 600 count for tw8

# current_root = 'C:/Users/伍嘉宏/Desktop/data/20200612'
# current_root = '/home/cq/Downloads/optical_encoder_compensation_2ms'
# current_root = 'C:/Users/伍嘉宏/Desktop/EncoderCompensationDataFit/optical_encoder_data_20200512/1'
# test_sucess_m = None
heidenhain_data = None
icmu_data = None
tw8_data = None
# popt_m = None
tw8_absolute_data = None
icmu_fit_result = None
tw8_fit_result = None


# 定义目标函数
def func1(theta, w, c, a1, b1):
    return (a1 * np.cos(w * theta) + b1 * np.sin(w * theta) + c)


def func2(theta, w, c, a1, b1, a2, b2):
    return (a1 * np.cos(w * theta) + b1 * np.sin(w * theta)
            + a2 * np.cos(2 * w * theta) + b2 * np.sin(2 * w * theta) + c)


def func3(theta, w, c, a1, b1, a2, b2, a3, b3):
    return (a1 * np.cos(w * theta) + b1 * np.sin(w * theta)
            + a2 * np.cos(2 * w * theta) + b2 * np.sin(2 * w * theta)
            + a3 * np.cos(3 * w * theta) + b3 * np.sin(3 * w * theta) + c)


def func4(theta, w, c, a1, b1, a2, b2, a3, b3, a4, b4):
    return (a1 * np.cos(w * theta) + b1 * np.sin(w * theta)
            + a2 * np.cos(2 * w * theta) + b2 * np.sin(2 * w * theta)
            + a3 * np.cos(3 * w * theta) + b3 * np.sin(3 * w * theta)
            + a4 * np.cos(4 * w * theta) + b4 * np.sin(4 * w * theta) + c)


def LoadFile(data_path):
    DATA = []
    with open(data_path, 'r') as f:
        r = f.readlines()
        for data in r[1:-1]:
            DATA.append([float(i) for i in data.split()])

    DATA = np.array(DATA)
    # print(DATA.shape)

    Isrtick = DATA[begin_row:end_row, isrtick_col]
    Encoder_Heidenhain = DATA[begin_row:end_row, heidenhain_col]
    Encoder_Icmu = DATA[begin_row:end_row, icmu_col]
    Encoder_tw8 = DATA[begin_row:end_row, tw8_col]
    PresentMode = DATA[begin_row:end_row, present_mode_col]
    IndexCurrent = DATA[begin_row:end_row, index_current_col]

    return Isrtick, Encoder_Heidenhain, Encoder_Icmu, Encoder_tw8, PresentMode, IndexCurrent


def read_sequence_file(path, extension='mode1_encoder_compensation.xls'):
    try:
        files = [f for f in os.listdir(path) if f.endswith(extension)]
    except FileNotFoundError as e:
        print(e)
    # sys.exit()
    # digital = sorted([int(re.sub("\D", "", f)) for f in files])
    # return ["log%d.xls" % i for i in digital]
    return files


def LoadData(current_root):
    # dir_char = sorted(os.listdir(current_root))
    dir_char = read_sequence_file(current_root, "mode1_encoder_compensation.xls")

    # Load Data and Combine Data
    isrtick_raw = []
    heidenhain_raw = []
    icmu_raw = []
    tw8_raw = []
    present_mode_raw = []
    index_current_raw = []
    # print(dir_char)

    for dir_i in dir_char:
        current_dir = current_root + '/' + dir_i
        isrtick_raw_temp, heidenhain_raw_temp, icmu_raw_temp, tw8_raw_temp, present_mode_temp, index_current_temp = LoadFile(
            current_dir)
        isrtick_raw.append(isrtick_raw_temp)
        heidenhain_raw.append(heidenhain_raw_temp)
        icmu_raw.append(icmu_raw_temp)
        tw8_raw.append(tw8_raw_temp)
        present_mode_raw.append(present_mode_temp)
        index_current_raw.append(index_current_temp)
    isrtick_raw = np.array(isrtick_raw)
    isrtick_raw = np.concatenate(isrtick_raw, axis=0)
    heidenhain_raw = np.array(heidenhain_raw)
    heidenhain_raw = np.concatenate(heidenhain_raw, axis=0)
    icmu_raw = np.array(icmu_raw)
    icmu_raw = np.concatenate(icmu_raw, axis=0)
    tw8_raw = np.array(tw8_raw)
    tw8_raw = np.concatenate(tw8_raw, axis=0)
    present_mode_raw = np.array(present_mode_raw)
    present_mode_raw = np.concatenate(present_mode_raw, axis=0)
    index_current_raw = np.array(index_current_raw)
    index_current_raw = np.concatenate(index_current_raw, axis=0)

    # Data Pick
    valid_data_start = 0
    for i in range(len(present_mode_raw)):
        if present_mode_raw[i] == 1 and index_current_raw[i] == 1:
            valid_data_start = i
            break
        else:
            continue

    # print(j)

    # plt.plot(tw8_raw, label='isr_raw')
    # plt.legend(loc=1)

    isrtick_pick = isrtick_raw[valid_data_start:end_row]
    heidenhain_pick = heidenhain_raw[valid_data_start:end_row]
    icmu_pick = icmu_raw[valid_data_start:end_row]
    tw8_pick = tw8_raw[valid_data_start:end_row]

    # plt.plot(heidenhain_raw, label='isr_pick')
    # plt.legend(loc=1)

    isrtick_filt = isrtick_pick
    heidenhain_filt = heidenhain_pick
    icmu_filt = icmu_pick
    tw8_filt = tw8_pick

    # Data Filter
    for i in range(len(isrtick_pick) - 1):
        if isrtick_pick[i + 1] - isrtick_pick[i] != 20:
            # isrtick_filt[i] = isrtick_pick[i-1] + 20
            heidenhain_filt[i] = heidenhain_pick[i - 1]
            icmu_filt[i] = icmu_pick[i - 1]
            tw8_filt[i] = tw8_pick[i - 1]
        else:
            continue

    # plt.plot(heidenhain_pick, label='isr_filt')
    # plt.legend(loc=1)

    # plt.show()

    return heidenhain_filt, icmu_filt, tw8_filt


def DataFit(ref_in, ref_max_bit, pos_in, pos_max_bit):
    fit_result = dict()

    # Data Processing2
    ref = ref_in / 2 ** ref_max_bit * 2 * np.pi
    pos = (1 - pos_in / 2 ** pos_max_bit) * 2 * np.pi

    err = ref - pos

    # plt.plot(ref, label='ref')
    # plt.plot(pos, label='pos')
    # plt.plot(err, label='err')
    # plt.legend(loc=1)

    # plt.show()

    xdata = pos
    ydata = err

    popt, pcov = curve_fit(func4, xdata, ydata)
    fit_result['popt'] = popt.tolist()

    encoder_err_rad_raw = ydata
    encoder_err_rad_fit = func4(xdata, *popt)
    encoder_err_rad_after_fit = abs(encoder_err_rad_fit - encoder_err_rad_raw)
    encoder_err_max_rad_after_fit = max(encoder_err_rad_after_fit)
    encoder_err_max_rad_raw = max(encoder_err_rad_raw)

    fit_result['max_encoder_err_pulse_raw'] = (encoder_err_max_rad_raw - popt[1]) / (2 * np.pi) * 2 ** pos_max_bit
    fit_result['max_encoder_err_pulse_fit'] = encoder_err_max_rad_after_fit / (2 * np.pi) * 2 ** pos_max_bit

    # print(fit_result['max_encoder_err_pulse_fit'])

    standard_count = standard_deg/360*(2**pos_max_bit);


    if fit_result['max_encoder_err_pulse_fit'] < standard_count:
        fit_result['test_successed'] = True
    else:
        fit_result['test_successed'] = False

    return fit_result


icmu_plot_elem = {
    'figure': 'Icmu Compensation Result',
    'title1': 'Heidehain_DATA-ICMU_DATA After Fit Or No',
    'title2': 'Heidehain_DATA-ICMU_DATA After Compensation Or No',
    'x_label': 'ICMU_DATA(rad)',
    'y_label': 'ERR(count)'
}
tw8_plot_elem = {
    'figure': 'Tw8 Compensation Result',
    'title1': 'Heidehain_DATA-TW8_DATA After Fit Or No',
    'title2': 'Heidehain_DATA-TW8_DATA After Compensation Or No',
    'x_label': 'TW8_DATA(rad)',
    'y_label': 'ERR(count)'
}


def DataPlot(ref_in, ref_max_bit, pos_in, pos_max_bit, popt, plot_elem):
    # if test_successed == 1:
    # Data Processing
    ref = ref_in / 2 ** ref_max_bit * 2 * np.pi
    # pos = (1 - pos_in / 2 ** 22) * 2 * np.pi #光
    pos = (1 - pos_in / 2 ** pos_max_bit) * 2 * np.pi  # 磁
    err = ref - pos

    xdata = pos
    ydata = err

    plt.figure(plot_elem['figure'])
    plt.subplot(2, 1, 1)
    plt.plot(xdata, ydata, 'b-', label='err_raw')
    plt.plot(xdata, func4(xdata, *popt), 'r-', label='err_fit')
    plt.xlabel(plot_elem['x_label'])
    plt.ylabel(plot_elem['y_label'])
    plt.legend(loc=1)
    plt.title(plot_elem['title1'])
    # plt.ylim(-4 * np.abs(popt[2]) + popt[1], 4 * np.abs(popt[2]) + popt[1])

    plt.subplot(2, 1, 2)
    plt.plot(xdata, ((ydata - func4(xdata, *popt)) * 2 ** pos_max_bit / (2 * np.pi)), 'r-', label='compensate_err')
    plt.plot(xdata, (ydata - popt[1]) * 2 ** pos_max_bit / (2 * np.pi), 'b-', label='no_compensate_err')
    plt.xlabel(plot_elem['x_label'])
    plt.ylabel(plot_elem['y_label'])
    plt.legend(loc=1)
    plt.title(plot_elem['title2'])
    # plt.ylim((-4 * np.abs(popt[2])) * 2 ** pos_max_bit / (2 * np.pi),
    #          (4 * np.abs(popt[2])) * 2 ** pos_max_bit / (2 * np.pi))
    plt.show()


# plt.show()

def Param_Plot(fit_result, title):
    # print(title)
    # print('test_sucessed:', fit_result['test_successed'])
    # print('max_encoder_err_pulse_raw:', fit_result['max_encoder_err_pulse_raw'])
    # print('max_encoder_err_pulse_fit:', fit_result['max_encoder_err_pulse_fit'])

    popt_list = ''
    param_name = [',//w', ',//c', ',//a1', ',//b1', ',//a2', ',//b2', ',//a3', ',//b3', ',//a4', ',//b4']
    # print('popt:')
    # print('//[', 'V2J2R005', ']')
    for i in range(0, 10):
        popt_list = popt_list + (str(fit_result['popt'][i]) + param_name[i] + '\n')
    # print(fit_result['popt'][i], param_name[i])

    # print('\n')

    return popt_list


def encoder_data():
    global heidenhain_data, icmu_data, tw8_data, tw8_absolute_data, icmu_fit_result, tw8_fit_result

    icmu_fit_result = dict()
    tw8_fit_result = dict()

    # Return Value#
    heidenhain_data, icmu_data, tw8_data = LoadData(current_root)

    # print('a',icmu_data[0])
    icmu_origin_raw_rad = (2 ** 20 * 2 - icmu_data[0]) * (2 * np.pi) / (2 ** 20)
    # print('b',icmu_origin_raw_rad)

    popt_tmp = icmu_fit_result['popt']
    popt_tmp[1] = 0
    icmu_origin_compensate_rad = icmu_origin_raw_rad + func4(icmu_origin_raw_rad, *popt_tmp)
    # print('c',icmu_origin_compensate_rad)

    icmu_origin_compensate_pulse = icmu_origin_compensate_rad * (2 ** 20) / (2 * np.pi)
    # print('d',icmu_origin_compensate_pulse)

    icmu_origin_compensate_pulse = icmu_origin_compensate_pulse % (2 ** 20)
    # print('e',icmu_origin_compensate_pulse)

    icmu_origin_compensate_pulse = 2 ** 20 - icmu_origin_compensate_pulse
    # print('f',icmu_origin_compensate_pulse)

    tw8_origin_compensate = icmu_origin_compensate_pulse / 2 ** 20 * 2 ** 22
    tw8_absolute_data = tw8_data + tw8_origin_compensate

    tw8_fit_result = DataFit(heidenhain_data, 25, tw8_absolute_data, 22)

    icmu_popt_list = Param_Plot(icmu_fit_result, 'icmu_param:')
    tw8_popt_list = Param_Plot(tw8_fit_result, 'tw8_param:')

    # Plot Data #
    return icmu_fit_result, tw8_fit_result, icmu_popt_list, tw8_popt_list


def encoderPlot():
    DataPlot(heidenhain_data, 25, icmu_data, 20, icmu_fit_result['popt'], icmu_plot_elem)
    DataPlot(heidenhain_data, 25, tw8_absolute_data, 22, tw8_fit_result['popt'], tw8_plot_elem)


def ui_icmu_fit_result(current_root):
    global heidenhain_data, icmu_data, tw8_data, tw8_absolute_data, icmu_fit_result, tw8_fit_result

    icmu_fit_result = dict()

    # Return Value#
    try:
        heidenhain_data, icmu_data, tw8_data = LoadData(current_root)

        icmu_fit_result = DataFit(heidenhain_data, 25, icmu_data, 20)

        # print('icmu_fit_result', icmu_fit_result)

        icmu_popt_list = Param_Plot(icmu_fit_result, 'icmu_param:')

        # Plot Data #
        return icmu_fit_result, icmu_popt_list
    except Exception as e:
        print('ui_icmu_fit_result')
        print(e)
    return None, None


def ui_tw8_fit_result():
    global heidenhain_data, icmu_data, tw8_data, tw8_absolute_data, icmu_fit_result, tw8_fit_result
    try:
        tw8_fit_result = dict()

        print('icmu_data[0]',icmu_data[0])
        icmu_origin_raw_rad = (2 ** 20 * 2 - icmu_data[0]) * (2 * np.pi) / (2 ** 20)
        print('icmu_origin_raw_rad',icmu_origin_raw_rad)

        popt_tmp = icmu_fit_result['popt'].copy()
        popt_tmp[1] = 0
        icmu_origin_compensate_rad = icmu_origin_raw_rad + func4(icmu_origin_raw_rad, *popt_tmp)
        print('icmu_origin_compensate_rad',icmu_origin_compensate_rad)

        icmu_origin_compensate_pulse = icmu_origin_compensate_rad * (2 ** 20) / (2 * np.pi)
        print('icmu_origin_compensate_pulse',icmu_origin_compensate_pulse)

        icmu_origin_compensate_pulse = icmu_origin_compensate_pulse % (2 ** 20)
        print('icmu_origin_compensate_pulse',icmu_origin_compensate_pulse)

        icmu_origin_compensate_pulse = 2 ** 20 - icmu_origin_compensate_pulse
        print('icmu_origin_compensate_pulse',icmu_origin_compensate_pulse)

        tw8_origin_icmu_compensate_pulse = icmu_origin_compensate_pulse / 2 ** 20 * 2 ** 22
        print('tw8_origin_compensate', tw8_origin_icmu_compensate_pulse)

        tw8_origin_tw8_pulse = tw8_data[0]
        tw8_absolute_data = tw8_data + tw8_origin_icmu_compensate_pulse - tw8_origin_tw8_pulse
        print('tw8_data', tw8_data[0])
        print('tw8_absolute_data', tw8_absolute_data[0])

        # plt.plot(tw8_data)

        tw8_fit_result = DataFit(heidenhain_data, 25, tw8_absolute_data, 22)

        tw8_popt_list = Param_Plot(tw8_fit_result, 'tw8_param:')

        # Plot Data #
        return tw8_fit_result, tw8_popt_list
    except Exception as e:
        print('ui_tw8_fit_result')
        print(e)
    return None, None


def ui_icmu_plot_result():
    try:
        DataPlot(heidenhain_data, 25, icmu_data, 20, icmu_fit_result['popt'], icmu_plot_elem)
    except Exception as e:
        print('ui_icmu_plot_result')
        print(e)


def ui_tw8_plot_result():
    try:
        DataPlot(heidenhain_data, 25, tw8_absolute_data, 22, tw8_fit_result['popt'], tw8_plot_elem)
    except Exception as e:
        print('ui_tw8_plot_result')
        print(e)


if __name__ == '__main__':
    # encoder_data()
    # encoderPlot()

    # 磁编码器补偿情况下：
    #   *点击‘数据处理’：调用 ui_icmu_fit_result()
    #   *点击‘画图’：调用 ui_icmu_plot_result()

    # 光编码器补偿情况下：
    #   *点击‘数据处理’：调用 ui_icmu_fit_result()，然后在调用 ui_tw8_fit_result()
    #   *点击磁编码器补偿框内的‘画图’：调用 ui_icmu_plot_result()
    #   *点击光编码器补偿框内的‘画图’：调用 ui_tw8_plot_result()

    # 注意：两种情况下的‘数据处理’按钮都用同一个就行，按马达ID来区分是哪种情况。
    #     磁和光编码器补偿的‘画图’按钮，设两个独立开来，能够独立查看不同情况下的画图

    icmu_fit_result, icmu_popt_list = ui_icmu_fit_result(current_root)
    ui_icmu_plot_result()
    tw8_fit_result, tw8_popt_list = ui_tw8_fit_result()
    ui_tw8_plot_result()
    print(icmu_fit_result, icmu_popt_list)
    print(tw8_fit_result, tw8_popt_list)

# plt.show()
# print(__name__)
